#!/usr/bin/env python3
"""
数据库初始化脚本
用于创建数据库表和插入初始数据
"""

import os
import sys
from sqlalchemy import create_engine, text
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

def init_database():
    """初始化数据库"""
    try:
        # 获取数据库连接信息
        database_url = os.getenv("DATABASE_URL")
        if not database_url:
            print("❌ 未找到数据库连接配置，请检查 .env 文件")
            return False
        
        # 创建数据库引擎
        engine = create_engine(database_url)
        
        print("🔄 开始初始化数据库...")
        
        # 读取并执行数据库创建脚本
        with open("create_database.sql", "r", encoding="utf-8") as f:
            create_sql = f.read()
        
        # 分割SQL语句并执行
        statements = [stmt.strip() for stmt in create_sql.split(';') if stmt.strip()]
        
        with engine.connect() as connection:
            for statement in statements:
                if statement:
                    connection.execute(text(statement))
                    connection.commit()
        
        print("✅ 数据库表创建成功")
        
        # 读取并执行数据初始化脚本
        with open("init_data.sql", "r", encoding="utf-8") as f:
            init_sql = f.read()
        
        # 分割SQL语句并执行
        statements = [stmt.strip() for stmt in init_sql.split(';') if stmt.strip()]
        
        with engine.connect() as connection:
            for statement in statements:
                if statement:
                    connection.execute(text(statement))
                    connection.commit()
        
        print("✅ 初始数据插入成功")
        
        # 验证数据
        with engine.connect() as connection:
            user_count = connection.execute(text("SELECT COUNT(*) FROM users")).scalar()
            book_count = connection.execute(text("SELECT COUNT(*) FROM books")).scalar()
            
            print(f"📊 数据统计:")
            print(f"   - 用户数量: {user_count}")
            print(f"   - 图书数量: {book_count}")
        
        print("🎉 数据库初始化完成！")
        return True
        
    except Exception as e:
        print(f"❌ 数据库初始化失败: {e}")
        return False

if __name__ == "__main__":
    success = init_database()
    sys.exit(0 if success else 1)